"""
对话记忆管理器

功能：
- 管理用户与系统的对话历史
- 支持多线程对话管理
- 提供对话摘要功能
- 支持数据持久化

作者：System
创建时间：2025-01-13
"""

from typing import List, Dict, Optional
from datetime import datetime
import json
import os
from pathlib import Path


class ConversationMemory:
    """对话记忆管理器"""
    
    def __init__(self, max_history: int = 10, storage_path: str = "cache/conversation_memory.json"):
        """
        初始化对话记忆
        
        Args:
            max_history: 保留的最大对话轮数（默认10轮）
            storage_path: 对话历史存储路径
        """
        self.max_history = max_history
        self.storage_path = storage_path
        self.conversations = {}  # {thread_id: [messages]}
        
        # 自动加载已有的对话历史
        self.load_from_file()
    
    def add_message(self, thread_id: str, role: str, content: str) -> None:
        """
        添加一条消息到对话历史
        
        Args:
            thread_id: 线程ID（用于区分不同的对话会话）
            role: 消息角色（"user" 或 "assistant"）
            content: 消息内容
        """
        if thread_id not in self.conversations:
            self.conversations[thread_id] = []
        
        message = {
            "role": role,  # "user" or "assistant"
            "content": content,
            "timestamp": datetime.now().isoformat()
        }
        
        self.conversations[thread_id].append(message)
        
        # 保持最大历史记录数（每轮对话包含用户和助手两条消息）
        if len(self.conversations[thread_id]) > self.max_history * 2:
            self.conversations[thread_id] = self.conversations[thread_id][-self.max_history * 2:]
        
        # 自动保存
        self.save_to_file()
    
    def get_recent_messages(self, thread_id: str, n: int = 5) -> List[Dict]:
        """
        获取最近的n轮对话
        
        Args:
            thread_id: 线程ID
            n: 获取的对话轮数（默认5轮）
            
        Returns:
            最近的消息列表
        """
        if thread_id not in self.conversations:
            return []
        
        # 每轮包含2条消息（用户+助手），所以返回 n*2 条
        return self.conversations[thread_id][-(n*2):]
    
    def get_all_messages(self, thread_id: str) -> List[Dict]:
        """
        获取某个线程的所有对话消息
        
        Args:
            thread_id: 线程ID
            
        Returns:
            所有消息列表
        """
        return self.conversations.get(thread_id, [])
    
    def get_conversation_summary(self, thread_id: str, max_length: int = 100) -> str:
        """
        获取对话摘要（用于上下文）
        
        Args:
            thread_id: 线程ID
            max_length: 每条消息的最大长度（默认100字符）
            
        Returns:
            格式化的对话摘要字符串
        """
        messages = self.get_recent_messages(thread_id)
        if not messages:
            return ""
        
        summary = "最近对话:\n"
        for msg in messages:
            role = "用户" if msg["role"] == "user" else "助手"
            # 截断过长的内容
            content = msg["content"]
            if len(content) > max_length:
                content = content[:max_length] + "..."
            summary += f"{role}: {content}\n"
        
        return summary
    
    def get_conversation_count(self, thread_id: str) -> int:
        """
        获取某个线程的对话消息总数
        
        Args:
            thread_id: 线程ID
            
        Returns:
            消息总数
        """
        return len(self.conversations.get(thread_id, []))
    
    def get_active_threads(self) -> List[str]:
        """
        获取所有活跃的线程ID列表
        
        Returns:
            线程ID列表
        """
        return list(self.conversations.keys())
    
    def clear_conversation(self, thread_id: str) -> None:
        """
        清空某个线程的对话历史
        
        Args:
            thread_id: 线程ID
        """
        if thread_id in self.conversations:
            del self.conversations[thread_id]
            self.save_to_file()
    
    def clear_all_conversations(self) -> None:
        """清空所有对话历史（慎用）"""
        self.conversations = {}
        self.save_to_file()
    
    def save_to_file(self) -> None:
        """保存对话历史到文件"""
        try:
            # 确保目录存在
            Path(self.storage_path).parent.mkdir(parents=True, exist_ok=True)
            
            with open(self.storage_path, 'w', encoding='utf-8') as f:
                json.dump(self.conversations, f, ensure_ascii=False, indent=2)
        except Exception as e:
            print(f"保存对话历史失败: {e}")
    
    def load_from_file(self) -> None:
        """从文件加载对话历史"""
        try:
            if os.path.exists(self.storage_path):
                with open(self.storage_path, 'r', encoding='utf-8') as f:
                    self.conversations = json.load(f)
        except FileNotFoundError:
            pass
        except Exception as e:
            print(f"加载对话历史失败: {e}")
    
    def get_statistics(self) -> Dict:
        """
        获取对话统计信息
        
        Returns:
            包含统计信息的字典
        """
        total_threads = len(self.conversations)
        total_messages = sum(len(msgs) for msgs in self.conversations.values())
        
        return {
            "total_threads": total_threads,
            "total_messages": total_messages,
            "active_threads": self.get_active_threads()
        }


# 创建全局单例（可选）
_global_conversation_memory = None

def get_conversation_memory(max_history: int = 10) -> ConversationMemory:
    """
    获取全局对话记忆管理器单例
    
    Args:
        max_history: 最大历史记录数
        
    Returns:
        ConversationMemory实例
    """
    global _global_conversation_memory
    if _global_conversation_memory is None:
        _global_conversation_memory = ConversationMemory(max_history=max_history)
    return _global_conversation_memory